import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import SGDRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# ==================== 初始化设置 ====================
plt.rcParams['font.sans-serif'] = ['KaiTi']  # 中文字体
plt.rcParams['mathtext.fontset'] = 'stix'    # 数学字体
plt.rcParams['axes.unicode_minus'] = False   # 负号显示

# ==================== 数据准备 ====================
np.random.seed(42)
X = 2 * np.random.rand(100, 1)
y = 4 + 3 * X + np.random.randn(100, 1)

# 划分训练测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 数据标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# ==================== SGD实现ElasticNet ====================
# 参数说明：
# penalty='elasticnet': 启用弹性网络正则化
# alpha: 总正则化强度
# l1_ratio: L1正则化比例(0=Ridge, 1=Lasso)
# max_iter: 最大迭代次数
# eta0: 初始学习率
# learning_rate: 学习率调整策略
sgd_reg = SGDRegressor(
    penalty='elasticnet',
    alpha=0.5,          # 总正则化强度
    l1_ratio=0.5,       # L1/L2混合比例
    max_iter=10000,     # 最大迭代次数
    eta0=0.01,          # 初始学习率
    learning_rate='constant',  # 固定学习率
    random_state=42     # 随机种子
)

# 训练模型(y需要转换为一维数组)
sgd_reg.fit(X_train_scaled, y_train.ravel())

# ==================== 模型评估 ====================
# 训练集预测
y_train_pred = sgd_reg.predict(X_train_scaled)
train_mse = mean_squared_error(y_train, y_train_pred)

# 测试集预测
y_test_pred = sgd_reg.predict(X_test_scaled)
test_mse = mean_squared_error(y_test, y_test_pred)

# ==================== 结果输出 ====================
print("===== SGD ElasticNet回归结果 =====")
print(f"截距项(w0): {sgd_reg.intercept_[0]:.4f}")
print(f"系数(w1): {sgd_reg.coef_[0]:.4f}")
print(f"训练集MSE: {train_mse:.4f}")
print(f"测试集MSE: {test_mse:.4f}")
print(f"非零系数数量: {np.sum(sgd_reg.coef_ != 0)}")

# ==================== 可视化 ====================
plt.figure(figsize=(12, 6))

# 绘制数据点
plt.scatter(X_train, y_train, color='blue', label='训练数据', alpha=0.6)
plt.scatter(X_test, y_test, color='green', label='测试数据', alpha=0.6)

# 绘制拟合线
X_plot = np.linspace(0, 2, 100).reshape(-1, 1)
X_plot_scaled = scaler.transform(X_plot)
y_plot = sgd_reg.predict(X_plot_scaled)
plt.plot(X_plot, y_plot, color='red', linewidth=2, label='SGD ElasticNet拟合线')

# 图表装饰
plt.title('SGD实现ElasticNet回归 (α=0.5, L1比例=0.5)', fontsize=14)
plt.xlabel('X特征值', fontsize=12)
plt.ylabel('y标签值', fontsize=12)
plt.legend(fontsize=10)
plt.grid(True)
plt.show()
